#include "rolling_shutter_correct.h"
#include "QueueUtil.h"

void test_rolling_shutter_correct()
{

}

/**
 *
 * @brief calculating weights for quaternion interpolation
 *
 * @param x init weight
 * @param easing smoothing method
 *
 */
float compute_interpolate_weight(float x, Easing easing)
{
    switch (easing) {
        case EaseIn:
            x = 1.0 - cosf(x * PI / 2.0);
            break;
        case EaseOut:
            x = sinf(x * PI / 2.0);
            break;
        case EaseInOut:
            x = -0.5 * (cosf(PI * x) - 1.0);
            break;
        default:
            break;
    }

    return x;
}

/**
 *
 * @brief get origin imu pose at time_stamp(image timestamp)
 *
 * @param imu_queue imu queue
 * @param imu_length imu length
 * @param time_stamp image timestamp
 *
 */

Quaternion org_quat_at_timestamp(Queue* imu_queue, uint32_t imu_length, int32_t time_stamp)
{
    // if (self.duration_ms <= 0.0)
    // {
    //     Quaternion q;
    //     setQuaternion(&q, 1.f, 0.f, 0.f, 0.f);
    //     return q;
    // }

    if (imu_length == 1)
    {
        Node* imu_front = imu_queue->front;
        if(imu_front->timeStamp == time_stamp) {
            return imu_front->q;
        }
        printf("Error0 org_quat_at_timestamp (imu_length == 1) \n");
        Quaternion q_wrong = {1.f,0.f,0.f,0.f};
        return q_wrong;
    }

    Node* firstElementCopy = getFirstElement(imu_queue);
    Node* lastElementCopy = getLastElement(imu_queue);

    if(!firstElementCopy || !lastElementCopy)
    {
        printf("Error org_quat_at_timestamp getFirstElement getLastElement NULL \n");
        Quaternion q_wrong = {1.f,0.f,0.f,0.f};
        return q_wrong;
    }

    int32_t first_ts = firstElementCopy->timeStamp;
    int32_t last_ts = lastElementCopy->timeStamp;
    int32_t lookup_ts = fmax(fmin(time_stamp, last_ts), first_ts);

    Node* current = imu_queue->front;
    Quaternion quat1 = current->q;
    float time_stamp1 = current->timeStamp;

    while (current != NULL) {

        if(current->timeStamp >= lookup_ts){

            if(current->timeStamp == lookup_ts) {
                return current->q;
            }

            time_stamp1 = time_stamp1 / 1000.f;
            float time_stamp2 = current->timeStamp / 1000.f;
            float time_stamp = lookup_ts / 1000.f;
            float time_delta = time_stamp2 - time_stamp1;
            float alpha = (time_stamp - time_stamp1) / time_delta;
            // Easing e = get_easing(eas1, eas2);
            Easing e = EaseInOut;
            float x = compute_interpolate_weight(alpha,e);

            Quaternion quat2 = current->q;
            Quaternion q_interpolate;
            slerp(&quat1, &quat2, x, &q_interpolate);
            return q_interpolate;
        }
        else{
            quat1 = current->q;
            time_stamp1 = current->timeStamp;
            current = current->next;
        }
    }

    printf("Error1 org_quat_at_timestamp \n");
    Quaternion q_wrong = {1.f,0.f,0.f,0.f};
    return q_wrong;
}

/**
 *
 * @brief get imu pose(before smoothed to smmoothed) at time_stamp(image timestamp)
 *
 * @param imu_queue imu queue
 * @param imu_length imu length
 * @param time_stamp image timestamp
 *
 */
Quaternion smoothed_quat_at_timestamp(Queue* imu_queue, uint32_t imu_length, int32_t time_stamp)
{
    // if (self.duration_ms <= 0.0)
    // {
    //     Quaternion q;
    //     setQuaternion(&q, 1.f, 0.f, 0.f, 0.f);
    //     return q;
    // }

    if (imu_length == 1)
    {
        Node* imu_front = imu_queue->front;
        if(imu_front->timeStamp == time_stamp) {
            return imu_front->diffQ;
        }
        printf("Error0 smoothed_quat_at_timestamp (imu_length == 1) \n");
        Quaternion q_wrong = {1.f,0.f,0.f,0.f};
        return q_wrong;
    }

    Node* firstElementCopy = getFirstElement(imu_queue);
    Node* lastElementCopy = getLastElement(imu_queue);
    if(!firstElementCopy || !lastElementCopy)
    {
        printf("Error smoothed_quat_at_timestamp getFirstElement getLastElement NULL \n");
        Quaternion q_wrong = {1.f,0.f,0.f,0.f};
        return q_wrong;
    }

    int32_t first_ts = firstElementCopy->timeStamp;
    int32_t last_ts = lastElementCopy->timeStamp;
    int32_t lookup_ts = fmax(fmin(time_stamp, last_ts), first_ts);

    Node* current = imu_queue->front;
    Quaternion quat1 = current->diffQ;
    float time_stamp1 = current->timeStamp;

    while (current != NULL) {

        if(current->timeStamp >= lookup_ts){

            if(current->timeStamp == lookup_ts) {
                return current->diffQ;
            }

            time_stamp1 = time_stamp1 / 1000.f;
            float time_stamp2 = current->timeStamp / 1000.f;
            float time_stamp = lookup_ts / 1000.f;
            float time_delta = time_stamp2 - time_stamp1;
            float alpha = (time_stamp - time_stamp1) / time_delta;
            // Easing e = get_easing(eas1, eas2);
            Easing e = EaseInOut;
            float x = compute_interpolate_weight(alpha,e);

            Quaternion quat2 = current->diffQ;
            Quaternion q_interpolate;
            slerp(&quat1, &quat2, x, &q_interpolate);
            return q_interpolate;
        }
        else{
            quat1 = current->diffQ;
            time_stamp1 = current->timeStamp;
            current = current->next;
        }
    }

    printf("Error1 smoothed_quat_at_timestamp \n");
    Quaternion q_wrong = {1.f,0.f,0.f,0.f};
    return q_wrong;
}


/**
 * TODO
 * @brief calculate new intrinsics based on cropping, scaling, etc. requirements
 *
 * @param camera_matrix intrinsics
 * @param fov
 *
 */
Matrix3 get_new_k(Matrix3* camera_matrix, float fov)
{
    // let horizontal_ratio = if params.lens.input_horizontal_stretch > 0.01 { params.lens.input_horizontal_stretch } else { 1.0 };
    // let img_dim_ratio = Self::get_ratio(params) / horizontal_ratio;
    // let out_dim = (params.output_width as f64, params.output_height as f64); // out_dim_x,out_dim_y

    float horizontal_ratio = 1.0f;
    float img_dim_ratio = 1.0f;
    uint32_t out_dim_x = 640;
    uint32_t out_dim_y = 1.0f;

    Matrix3 new_k;
    memcpy(&new_k, camera_matrix, sizeof(new_k));
    new_k.data[0][0] = camera_matrix->data[0][0] * img_dim_ratio / fov;
    new_k.data[1][1] = camera_matrix->data[1][1] * img_dim_ratio / fov;
    new_k.data[0][2] = (float)out_dim_x / 2.0f;
    new_k.data[1][2] = (float)out_dim_y / 2.0f;

    return new_k;
}

/**
 *
 * @brief compute the correction matrix for the row of the image
 *        the correction matrix is to correct the row of images to the img->time_stamp moment
 *
 * @param imu_queue the imu queue corresponding to the current image
 * @param img original image (shaking + rolling shutter), uniformly corrected to img->time_stamp moment
 * @param intrinsics intrinsics
 * @param extrinsics_cam_imu extrinsics_cam_imu
 * @param extrinsics_imu_cam extrinsics_imu_cam
 * @param row_time_stamp timestamp of the current row of image
 *
 */
Matrix3 q_compute_correct_rolling_shutter_matrix(Queue* imu_queue,
                                                 int img_timestamp,
                                                 Matrix3* intrinsics,   // 配置参数赋值
                                                 Matrix3* intrinsics_inv,
                                                 Matrix3* extrinsics_cam_imu,   // e_i_c的转置
                                                 Matrix3* extrinsics_imu_cam,   //配置参数赋值
                                                 int32_t row_time_stamp)
{
    Matrix3 correct_matrix;
    initUnitCircle(&correct_matrix);

    Matrix3 new_k = *intrinsics;
    Matrix3 new_k_inv = *intrinsics_inv;

    uint32_t imu_length = imu_queue->length;

    Matrix3 intrinsics_extrinsics = matrix_multiply(&new_k,extrinsics_cam_imu);

    //quaternion : smooth-before to smooth-after
    Quaternion q_imgsmooth_img = smoothed_quat_at_timestamp(imu_queue, imu_length,img_timestamp);
    
    Matrix3 mat_imgsmooth_img = toRotationMatrix(&q_imgsmooth_img);
    Matrix3 correct_matrix_temp = matrix_multiply(&intrinsics_extrinsics, &mat_imgsmooth_img);
    Matrix3 intrinsics_extrinsics_inv = matrix_multiply(extrinsics_imu_cam,&new_k_inv);
    correct_matrix = matrix_multiply(&correct_matrix_temp,&intrinsics_extrinsics_inv);

    return correct_matrix;
}

